import sys 
import gzip
import datetime
import threading
import math

# This program is for gov2 dataset
# updated by Wei 2012/11/26 evening(mostly rerun the whole thing again).
# I think this is good program to re-use for building easy structure for gov2 data as well.
# Let's measure the time to take as well.

# for production mode
# my default setting of the parameters will be 4 0/1/2/3 4 100000 using 4 processes, 4 threads each and make the last argument to be large like 10000
# the detail commands will be the following:
    # python programName 4 0 4 100000
    # python programName 4 1 4 100000
    # python programName 4 2 4 100000
    # python programName 4 3 4 100000
    # The format: python 20121116-gov2-scriptForBuildingEasyStructures.py 4 0 4 100000

# for debugging mode
# my default setting of the parameters will be 1 0 4 1 using 1 processes, 1 threads each and make the last argument to be 1 and we can see what happened during debug time


class ThreadClass(threading.Thread):
    # in production mode
    # STOP_THRESHOLD = -1
    
    # in debugging mode
    
    
    NUMBER_OF_THREADS = -1
    processID = -1
    totalNumberOfProcesses = -1
    targetSegmentFileNameList = []
    existingDocIDWARCTRECIDPairDict = {}
    def run(self):
        print "Process",self.processID,self.getName(),"Created!","Each thread need to process at most",STOP_THRESHOLD,'compressed files.'
        
        
        threadSpecificTargetSegmentFileNameList = []
        numberOfSegmentsForEachThread = math.ceil( len(self.targetSegmentFileNameList)/float(self.NUMBER_OF_THREADS) )
        threadID = int( self.getName().strip().split("-")[1] )
        lowerBoundForSpecifcThread = int( (threadID-1)*numberOfSegmentsForEachThread )
        upperBoundForSpecificThread = int( lowerBoundForSpecifcThread + numberOfSegmentsForEachThread )
        for currentTargetSegmentFileName in self.targetSegmentFileNameList[lowerBoundForSpecifcThread:upperBoundForSpecificThread]:
            threadSpecificTargetSegmentFileNameList.append(currentTargetSegmentFileName)
        
        #print self.targetSegmentFileNameList
        #print threadSpecificTargetSegmentFileNameList
        
        compressedFileCounter = -1
        for targetSegmentFileName in threadSpecificTargetSegmentFileNameList:
            
            if compressedFileCounter == self.STOP_THRESHOLD:
                print "Stopping","Process",self.processID,self.getName(),"now!"
                break
            
            numberOfDocumentsAlreadyProcessedInThisThread = 0
            inputFileName = "/data1/team/obukai/machine-learning-project-related/auxFiles/gov2DataPaths/" + targetSegmentFileName
            inputFileHandler = open(inputFileName,'r')
            
            outputFileNameForDoc = "/data5/team/obukai/the_new_trip_of_feature_generation/auxFilesALL/gov2DataDocsMapping/" + targetSegmentFileName
            outputFileNameForDocInCompressedForm = "/data5/team/obukai/the_new_trip_of_feature_generation/auxFilesALL/gov2DataDocsMappingInCompressedForm/" + targetSegmentFileName
            outputFileNameForCompressedFile = "/data5/team/obukai/the_new_trip_of_feature_generation/auxFilesALL/gov2DataCompressedFileMapping/" + targetSegmentFileName
            
            outputFileForDocHandler = open(outputFileNameForDoc, "w")
            outputFileForDocCompressedFileHandler = gzip.GzipFile( outputFileNameForDocInCompressedForm + ".gz", "wb")
            
            outputFileForCompressedFileHandler = open(outputFileNameForCompressedFile, "w")
            
            #outputFileForCompressedFileHandler.write("Hello" + "\n")
            
            
            

            
            for currentLine in inputFileHandler.readlines():
                inputProcessingDataFileName = currentLine.strip()
                compressedFileCounter += 1
                
                if compressedFileCounter == -1:
                    print "Skipping.. ",inputProcessingDataFileName
                    continue
                elif compressedFileCounter == self.STOP_THRESHOLD:
                    
                    break
                else:
                    print compressedFileCounter,"Process",self.processID,self.getName(),"Processing ",inputProcessingDataFileName
                    outputFileForCompressedFileHandler.write(str(compressedFileCounter) + " " + inputProcessingDataFileName + "\n")
                
                inputProcessingDataFileHandler = gzip.GzipFile(inputProcessingDataFileName,'r')
                #inputProcessingDataFileHandler.readline().strip().split(":")[0] == "WARC-TREC-ID" or 
                
                
                absoluteDocBeginningPositionInOneCompressFile = 0
                absoluteDocEndingPositionInOneCompressFile = 0
                currentDataLine = inputProcessingDataFileHandler.readline()
                
                # just for the beginning test
                # print "absolute beginning data line:",currentDataLine
                
                # the actual webpage content, many times
                flag1 = False
                while currentDataLine:
                    # step1: have the DOCNO extracted from the file
                    if currentDataLine.strip().startswith("<DOCNO>"):
                        currentWARCTRECID = currentDataLine.strip().split(">")[1].split("</")[0]
                        # print currentWARCTRECID
                        if currentWARCTRECID in self.existingDocIDWARCTRECIDPairDict:
                            flag1 = True
                            currentDocID = self.existingDocIDWARCTRECIDPairDict[currentWARCTRECID]
                        else:
                            flag1 = False
                        
                    
                    # step2: let's find out the beginning position of the individual document                                      
                    elif currentDataLine.strip().startswith("</DOCHDR>"):
                        absoluteDocBeginningPositionInOneCompressFile = inputProcessingDataFileHandler.tell()
                        
                    # step3: let's find out the ending position of the individual document
                    elif currentDataLine.strip().startswith("</DOC>"):
                        # why is 56, cause the </DOC> has 6 chars and the new line /n has one char, so I assume there are 7 chars and there are 56bytes
                        absoluteDocEndingPositionInOneCompressFile = inputProcessingDataFileHandler.tell() - 7
                        if flag1:
                            # print "currentWARCTRECID:",currentWARCTRECID
                            # print "absoluteDocBeginningPositionInOneCompressFile:",absoluteDocBeginningPositionInOneCompressFile
                            # print "absoluteDocEndingPositionInOneCompressFile:",absoluteDocEndingPositionInOneCompressFile
                            # print 
                            numberOfDocumentsAlreadyProcessedInThisThread += 1
                            # print "numberOfDocumentsAlreadyProcessedInThisThread:",numberOfDocumentsAlreadyProcessedInThisThread
                            outputFileForDocHandler.write(currentDocID + " " + currentWARCTRECID + " " + str(compressedFileCounter) + " " + str(absoluteDocBeginningPositionInOneCompressFile) + " " + str(absoluteDocEndingPositionInOneCompressFile) + "\n")
                            outputFileForDocCompressedFileHandler.write(currentDocID + " " + currentWARCTRECID + " " + str(compressedFileCounter) + " " + str(absoluteDocBeginningPositionInOneCompressFile) + " " + str(absoluteDocEndingPositionInOneCompressFile) + "\n")
                    currentDataLine = inputProcessingDataFileHandler.readline()
                    
                inputProcessingDataFileHandler.close()
                print "Total document recoreded:",numberOfDocumentsAlreadyProcessedInThisThread
                #flush these 3 files once the individual compressed file is finished.
                outputFileForDocHandler.flush()
                outputFileForDocCompressedFileHandler.flush()
                outputFileForCompressedFileHandler.flush()
                
            outputFileForDocHandler.close()
            outputFileForDocCompressedFileHandler.close()
            outputFileForCompressedFileHandler.close()


def loadTrecIDWhichNeedToBuildEasyAccessTo():
    # This method is currently being used.
    existingDocIDWARCTRECIDPairDict = {}
    inputFileName = "/data1/team/obukai/machine-learning-project-related/trec-related/trecCombine04_06/priorityTrecDocumentsIDs_include_approximating_BM25_AND_human_judged.txt"
    inputFileHandler = open(inputFileName, "r")
    passCounter = 0
        
    for currentLine in inputFileHandler.readlines():
        if currentLine.strip().split(" ")[1] not in existingDocIDWARCTRECIDPairDict:
            # Just put the TrecID index number into the dict instead of actual docID number or sth.
            existingDocIDWARCTRECIDPairDict[ currentLine.strip().split(" ")[1] ] =  currentLine.strip().split(" ")[0]
        else:
            passCounter += 1
        
    #print "len(existingDocIDWARCTRECIDPairDict)",len(existingDocIDWARCTRECIDPairDict)
    #print "passCounter",passCounter
    inputFileHandler.close()
    return existingDocIDWARCTRECIDPairDict

def loadTargetsSegmentNameTable():
    #step0: read the segment names.
    segmentNameList = []
    inputFileNameForThreadPathMapping = "/data1/team/obukai/machine-learning-project-related/auxFiles/threadPathMappingTableForGov2"
    inputFileHandlerForThreadPathMapping = open(inputFileNameForThreadPathMapping, "r")
    for currentLine in inputFileHandlerForThreadPathMapping.readlines():
        segmentNameList.append( currentLine.strip() )
    inputFileHandlerForThreadPathMapping.close()
    return segmentNameList

if len(sys.argv) < 5 :
    print "Usage: python scriptForBuildingEasyStructureForGov2Data.py <total # of processes> <current process ID> <# of threads> <# of compressed files needed to process>"
    sys.exit()
elif len(sys.argv) == 5:
    STOP_THRESHOLD = int(sys.argv[4])
    NUMBER_OF_THREADS = int(sys.argv[3])
    
    segmentNameList = loadTargetsSegmentNameTable()
    
    numberOfSegmentsForEachProcess = math.ceil( len(segmentNameList)/float(sys.argv[1]) )
    #print "numberOfSegmentsForEachProcess:",numberOfSegmentsForEachProcess

    #working old version
    
    for i in range(NUMBER_OF_THREADS):
        t = ThreadClass()
        t.STOP_THRESHOLD = STOP_THRESHOLD
        t.NUMBER_OF_THREADS = NUMBER_OF_THREADS
        t.totalNumberOfProcesses = int(sys.argv[1])
        t.processID = int(sys.argv[2])
        
        # Wei: 2012/10/05: The loadDocIDWARCTRECID1MPair() needed to be replaced.
        # The new replacement will be loadTrecIDWhichNeedToBuildEasyAccessTo()
        t.existingDocIDWARCTRECIDPairDict = loadTrecIDWhichNeedToBuildEasyAccessTo()
        
        #Let's assign the segment names based on the processID
        lowerBound = int( t.processID*numberOfSegmentsForEachProcess )
        upperBound = int( lowerBound + numberOfSegmentsForEachProcess )
        #print "lowerBound:",lowerBound
        #print "upperBound:",upperBound
        for currentSegmentName in segmentNameList[lowerBound:upperBound]:
            if currentSegmentName not in t.targetSegmentFileNameList:
                t.targetSegmentFileNameList.append(currentSegmentName)
            else:
                pass
        t.start()    
    #print "Run the Process", t.processID,"in",NUMBER_OF_THREADS,"Threads Done!"
else:
    print "Unknown Command"
    sys.exit() 


